package network

import (
	"context"
	"fmt"
	"io"
	"net"
	"sokit-go/utils"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

// Transfer 表示一个端口转发器
type Transfer struct {
	listener    net.Listener
	localAddr   string
	targetAddr  string
	protocol    string
	isRunning   bool
	mutex       sync.RWMutex
	connections map[net.Conn]bool
	connMutex   sync.RWMutex

	// 统计信息
	bytesTransferred  int64
	startTime         time.Time
	activeConnections int64

	// 性能优化
	config             *utils.Config
	bufferPool         sync.Pool
	ctx                context.Context
	cancel             context.CancelFunc
	maxConnections     int64
	currentConnections int64

	// 回调函数
	OnStarted         func()
	OnStopped         func()
	OnDataTransferred func(bytes int64)
	OnError           func(err error)
}

// NewTransfer 创建新的端口转发器实例
func NewTransfer() *Transfer {
	config := utils.GetConfig()
	transferConfig := config.GetTransferConfig()

	ctx, cancel := context.WithCancel(context.Background())

	t := &Transfer{
		isRunning:      false,
		connections:    make(map[net.Conn]bool),
		config:         config,
		ctx:            ctx,
		cancel:         cancel,
		maxConnections: int64(transferConfig.MaxConnections),
	}

	// 初始化缓冲区池
	t.bufferPool = sync.Pool{
		New: func() interface{} {
			return make([]byte, transferConfig.BufferSize)
		},
	}

	return t
}

// Start 开始端口转发
func (t *Transfer) Start(localHost, localPort, targetHost, targetPort, protocol string) error {
	t.mutex.Lock()
	defer t.mutex.Unlock()

	if t.isRunning {
		return fmt.Errorf("端口转发已在运行")
	}

	// 重新创建上下文，确保每次启动都有新的上下文
	t.ctx, t.cancel = context.WithCancel(context.Background())

	t.localAddr = fmt.Sprintf("%s:%s", localHost, localPort)
	t.targetAddr = fmt.Sprintf("%s:%s", targetHost, targetPort)
	t.protocol = strings.ToUpper(protocol)

	// 添加调试日志
	utils.GetLogger().Info("启动转发: %s -> %s (%s)", t.localAddr, t.targetAddr, t.protocol)

	var err error

	switch t.protocol {
	case "TCP":
		err = t.startTCPForward()
	case "UDP":
		err = t.startUDPForward()
	default:
		return fmt.Errorf("不支持的协议: %s", protocol)
	}

	if err != nil {
		utils.GetLogger().Error("转发启动失败: %v", err)
		if t.OnError != nil {
			t.OnError(err)
		}
		return err
	}

	t.isRunning = true
	t.startTime = time.Now()
	t.bytesTransferred = 0
	t.activeConnections = 0

	utils.GetLogger().Info("转发启动成功: %s -> %s", t.localAddr, t.targetAddr)

	if t.OnStarted != nil {
		t.OnStarted()
	}

	return nil
}

// startTCPForward 启动TCP端口转发
func (t *Transfer) startTCPForward() error {
	utils.GetLogger().Info("启动TCP转发监听: %s", t.localAddr)

	listener, err := net.Listen("tcp", t.localAddr)
	if err != nil {
		utils.GetLogger().Error("TCP监听失败: %v", err)
		return fmt.Errorf("监听本地端口失败: %v", err)
	}

	t.listener = listener
	utils.GetLogger().Info("TCP监听器创建成功: %s", t.localAddr)

	// 启动连接接受协程
	go t.acceptConnections()

	return nil
}

// acceptConnections 接受客户端连接
func (t *Transfer) acceptConnections() {
	defer func() {
		if t.listener != nil {
			t.listener.Close()
		}
	}()

	utils.GetLogger().Info("开始接受TCP连接: %s", t.localAddr)

	for {
		select {
		case <-t.ctx.Done():
			utils.GetLogger().Info("TCP接受协程被取消")
			return
		default:
		}

		if !t.isRunning {
			utils.GetLogger().Info("转发已停止，退出接受协程")
			break
		}

		conn, err := t.listener.Accept()
		if err != nil {
			if t.OnError != nil && t.isRunning {
				t.OnError(fmt.Errorf("接受连接失败: %v", err))
			}
			continue
		}

		utils.GetLogger().Info("接受新连接: %s -> %s", conn.RemoteAddr(), t.targetAddr)

		// 检查连接数限制
		if atomic.LoadInt64(&t.currentConnections) >= t.maxConnections {
			conn.Close()
			utils.GetLogger().Warn("连接数已达到最大值: %d", t.maxConnections)
			if t.OnError != nil {
				t.OnError(fmt.Errorf("连接数已达到最大值: %d", t.maxConnections))
			}
			continue
		}

		// 记录连接
		t.connMutex.Lock()
		t.connections[conn] = true
		atomic.AddInt64(&t.activeConnections, 1)
		t.connMutex.Unlock()

		// 处理连接
		go t.handleConnection(conn)
	}
}

// handleConnection 处理单个客户端连接
func (t *Transfer) handleConnection(clientConn net.Conn) {
	defer func() {
		clientConn.Close()
		t.connMutex.Lock()
		delete(t.connections, clientConn)
		atomic.AddInt64(&t.activeConnections, -1)
		atomic.AddInt64(&t.currentConnections, -1)
		t.connMutex.Unlock()
	}()

	atomic.AddInt64(&t.currentConnections, 1)

	// 连接到目标服务器
	transferConfig := t.config.GetTransferConfig()
	dialer := &net.Dialer{
		Timeout: time.Duration(transferConfig.WriteTimeout) * time.Second,
	}

	targetConn, err := dialer.DialContext(t.ctx, "tcp", t.targetAddr)
	if err != nil {
		if t.OnError != nil {
			t.OnError(fmt.Errorf("连接目标服务器失败: %v", err))
		}
		return
	}
	defer targetConn.Close()

	// 启动双向数据转发
	errChan := make(chan error, 2)

	go func() {
		err := t.forwardData(clientConn, targetConn, "client->target")
		errChan <- err
	}()

	go func() {
		err := t.forwardData(targetConn, clientConn, "target->client")
		errChan <- err
	}()

	// 等待任一方向的连接关闭
	select {
	case <-errChan:
	case <-t.ctx.Done():
	}
}

// forwardData 转发数据
func (t *Transfer) forwardData(src, dst net.Conn, direction string) error {
	defer dst.Close()

	// 从缓冲区池获取缓冲区
	buffer := t.bufferPool.Get().([]byte)
	defer t.bufferPool.Put(buffer)

	transferConfig := t.config.GetTransferConfig()
	readTimeout := time.Duration(transferConfig.ReadTimeout) * time.Second
	writeTimeout := time.Duration(transferConfig.WriteTimeout) * time.Second

	for {
		select {
		case <-t.ctx.Done():
			return nil
		default:
		}

		// 设置读取超时
		src.SetReadDeadline(time.Now().Add(readTimeout))

		n, err := src.Read(buffer)
		if err != nil {
			if err == io.EOF {
				return nil
			}
			if t.OnError != nil && t.isRunning {
				t.OnError(fmt.Errorf("读取数据失败 [%s]: %v", direction, err))
			}
			return err
		}

		if n > 0 {
			// 设置写入超时
			dst.SetWriteDeadline(time.Now().Add(writeTimeout))

			_, err = dst.Write(buffer[:n])
			if err != nil {
				if t.OnError != nil && t.isRunning {
					t.OnError(fmt.Errorf("写入数据失败 [%s]: %v", direction, err))
				}
				return err
			}

			// 更新统计信息（使用原子操作）
			atomic.AddInt64(&t.bytesTransferred, int64(n))

			if t.OnDataTransferred != nil {
				t.OnDataTransferred(int64(n))
			}
		}
	}
}

// startUDPForward 启动UDP端口转发
func (t *Transfer) startUDPForward() error {
	// 解析本地和目标地址
	localAddr, err := net.ResolveUDPAddr("udp", t.localAddr)
	if err != nil {
		return fmt.Errorf("解析本地地址失败: %v", err)
	}

	targetAddr, err := net.ResolveUDPAddr("udp", t.targetAddr)
	if err != nil {
		return fmt.Errorf("解析目标地址失败: %v", err)
	}

	// 创建UDP监听器
	listener, err := net.ListenUDP("udp", localAddr)
	if err != nil {
		return fmt.Errorf("创建UDP监听器失败: %v", err)
	}

	// 客户端连接映射（IP:Port -> 连接到目标服务器的连接）
	clientConnections := make(map[string]*net.UDPConn)
	connectionMutex := sync.RWMutex{}

	// 清理连接的函数
	cleanupConnection := func(clientKey string) {
		connectionMutex.Lock()
		defer connectionMutex.Unlock()
		if conn, exists := clientConnections[clientKey]; exists {
			conn.Close()
			delete(clientConnections, clientKey)
		}
	}

	// 启动监听协程
	go func() {
		buffer := make([]byte, 4096)
		for {
			if !t.IsRunning() {
				break
			}

			// 从客户端接收数据
			n, clientAddr, err := listener.ReadFromUDP(buffer)
			if err != nil {
				if !t.IsRunning() {
					break
				}
				if t.OnError != nil {
					t.OnError(fmt.Errorf("UDP读取错误: %v", err))
				}
				continue
			}

			clientKey := clientAddr.String()

			// 获取或创建到目标服务器的连接
			connectionMutex.RLock()
			targetConn, exists := clientConnections[clientKey]
			connectionMutex.RUnlock()

			if !exists {
				// 创建新的连接到目标服务器
				newConn, err := net.DialUDP("udp", nil, targetAddr)
				if err != nil {
					if t.OnError != nil {
						t.OnError(fmt.Errorf("连接目标服务器失败: %v", err))
					}
					continue
				}

				connectionMutex.Lock()
				clientConnections[clientKey] = newConn
				connectionMutex.Unlock()

				targetConn = newConn

				// 启动从目标服务器返回数据的协程
				go func(conn *net.UDPConn, client *net.UDPAddr, key string) {
					defer cleanupConnection(key)

					responseBuffer := make([]byte, 4096)
					// 设置超时时间，5分钟无数据则关闭连接
					conn.SetReadDeadline(time.Now().Add(5 * time.Minute))

					for {
						if !t.IsRunning() {
							break
						}

						n, err := conn.Read(responseBuffer)
						if err != nil {
							if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
								// 超时，清理连接
								break
							}
							if !t.IsRunning() {
								break
							}
							break
						}

						// 发送响应回客户端
						_, err = listener.WriteToUDP(responseBuffer[:n], client)
						if err != nil {
							if !t.IsRunning() {
								break
							}
							if t.OnError != nil {
								t.OnError(fmt.Errorf("UDP写入错误: %v", err))
							}
							break
						}

						// 更新统计信息（使用原子操作）
						atomic.AddInt64(&t.bytesTransferred, int64(n))

						if t.OnDataTransferred != nil {
							t.OnDataTransferred(int64(n))
						}

						// 重置超时时间
						conn.SetReadDeadline(time.Now().Add(5 * time.Minute))
					}
				}(targetConn, clientAddr, clientKey)
			}

			// 转发数据到目标服务器
			_, err = targetConn.Write(buffer[:n])
			if err != nil {
				if t.OnError != nil {
					t.OnError(fmt.Errorf("转发数据到目标服务器失败: %v", err))
				}
				cleanupConnection(clientKey)
				continue
			}

			// 更新统计信息（使用原子操作）
			atomic.AddInt64(&t.bytesTransferred, int64(n))

			if t.OnDataTransferred != nil {
				t.OnDataTransferred(int64(n))
			}
		}

		// 清理所有连接
		connectionMutex.Lock()
		for key, conn := range clientConnections {
			conn.Close()
			delete(clientConnections, key)
		}
		connectionMutex.Unlock()

		listener.Close()
	}()

	return nil
}

// Stop 停止端口转发
func (t *Transfer) Stop() error {
	t.mutex.Lock()
	defer t.mutex.Unlock()

	if !t.isRunning {
		return nil
	}

	t.isRunning = false

	// 取消所有操作
	t.cancel()

	// 关闭监听器
	if t.listener != nil {
		t.listener.Close()
		t.listener = nil
	}

	// 关闭所有活动连接
	t.connMutex.Lock()
	for conn := range t.connections {
		conn.Close()
	}
	t.connections = make(map[net.Conn]bool)
	atomic.StoreInt64(&t.activeConnections, 0)
	atomic.StoreInt64(&t.currentConnections, 0)
	t.connMutex.Unlock()

	if t.OnStopped != nil {
		t.OnStopped()
	}

	return nil
}

// GetActiveConnections 获取活动连接数
func (t *Transfer) GetActiveConnections() int {
	return int(atomic.LoadInt64(&t.activeConnections))
}

// GetBytesTransferred 获取传输字节数
func (t *Transfer) GetBytesTransferred() int64 {
	return atomic.LoadInt64(&t.bytesTransferred)
}

// IsRunning 检查传输是否运行
func (t *Transfer) IsRunning() bool {
	t.mutex.RLock()
	defer t.mutex.RUnlock()
	return t.isRunning
}

// GetTransferInfo 获取转发信息
func (t *Transfer) GetTransferInfo() (string, string, string, int64, time.Duration) {
	t.mutex.RLock()
	defer t.mutex.RUnlock()

	var duration time.Duration
	if !t.startTime.IsZero() {
		duration = time.Since(t.startTime)
	}

	return t.localAddr, t.targetAddr, t.protocol, t.bytesTransferred, duration
}

// GetTransferStats 获取传输统计信息
func (t *Transfer) GetTransferStats() (int64, time.Duration, float64) {
	t.mutex.RLock()
	startTime := t.startTime
	t.mutex.RUnlock()

	bytesTransferred := atomic.LoadInt64(&t.bytesTransferred)

	var duration time.Duration
	if !startTime.IsZero() {
		duration = time.Since(startTime)
	}

	var speed float64
	if duration.Seconds() > 0 {
		speed = float64(bytesTransferred) / duration.Seconds()
	}

	return bytesTransferred, duration, speed
}
